from model_classes.ort_loader import ORTLoader

import numpy  as np
import soundfile as sf
# from transformers import Wav2Vec2Processor

model_path = "./models/data2vec_onnx/model.onnx"

data_path = "./datasets/text2audio.wav"

if __name__ == "__main__":
    # 读取音频文件
    audio, sp = sf.read(data_path)
    print(f"音频采样率: {sp}Hz, 长度: {len(audio)} 个样本")

    # 初始化模型加载器
    ort_loader = ORTLoader(model_path)

    # 准备输入数据
    input_data = np.expand_dims(audio, axis=0).astype(np.float16)
    print(f"输入数据形状: {input_data.shape}")
    
    # process = Wav2Vec2Processor.from_pretrained("./models/data2vec_onnx/")
    # result = ort_loader.ExecInfer_v2(input_data, sp, 16000)

    # 执行推理和解码
    result = ort_loader.ExecInfer(input_data, sp, 16000)

    # 打印结果
    print("\n=== ASR 解码结果 ===")
    print(f"识别文本: '{str.lower(result['text'])}'")
    print(f"解码的token IDs: {result['decoded_tokens']}")
    print(f"Logits 形状: {result['logits'].shape}")
    
    # 显示每个解码的token
    print("\n=== Token 详情 ===")
    for i, token_id in enumerate(result['decoded_tokens']):
        token = ort_loader.id_to_token[token_id]
        print(f"Token {i}: ID={token_id}, 内容='{token}'")
    
    # 可选：显示前几个时间步的概率分布
    print("\n=== 前5个时间步的Top-3预测 ===")
    for t in range(min(5, len(result['logits']))):
        top3_ids = np.argsort(result['logits'][t])[-3:][::-1]
        top3_probs = np.sort(result['logits'][t])[-3:][::-1]
        print(f"时间步 {t}:")
        for i, (token_id, prob) in enumerate(zip(top3_ids, top3_probs)):
            token = ort_loader.id_to_token[token_id]
            print(f"  Top-{i+1}: '{token}' (ID={token_id}, logit={prob:.3f})")